/**
 * @file all_gather_matmul_custom.h
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include "kernel_operator.h"
#include "kernel_operator_intf.h"
#include "lib/matmul_intf.h"
#include "gather_mm.h"
#include "all_gather_matmul_custom_tiling.h"

using namespace AscendC;

template <class T>
struct BiasType {
    using type = float;
};
template <>
struct BiasType<half> {
    using type = half;
};

extern "C" __global__ __aicore__ void all_gather_matmul_custom(GM_ADDR aGM, GM_ADDR bGM, GM_ADDR biasGM, GM_ADDR cGM,
    GM_ADDR gatherOutGM, GM_ADDR workspaceGM, GM_ADDR tilingGM)
{
    if ASCEND_IS_AIV {
        return;
    }
    REGISTER_TILING_DEFAULT(AllGatherMatmulCustomTilingData);
    auto tiling = (__gm__ AllGatherMatmulCustomTilingData*)tilingGM;
    __gm__ void *mc2InitTiling = (__gm__ void *)(&(tiling->mc2InitTiling));
    __gm__ void *mc2CcTiling = (__gm__ void *)(&(tiling->mc2CcTiling));

    GET_TILING_DATA(tilingData, tilingGM);
    auto &&cfg         = tilingData.param;
    auto &&localTiling = tilingData.localTiling;
    auto &&tileTiling  = tilingData.tileTiling;
    auto &&tailTiling  = tilingData.tailTiling;
    const auto tileNum = cfg.tileNum;
    const auto tailNum = cfg.tailNum;

    const auto aTileCnt    = tileTiling.M * tileTiling.Ka;
    const auto aTileOffset = tileTiling.M * tileTiling.Ka * sizeof(A_DTYPE);
    const auto cTileOffset = tileTiling.M * tileTiling.N * sizeof(C_DTYPE);
    const auto aTailCnt    = tailTiling.M * tailTiling.Ka;
    const auto aRankCnt    = cfg.rankM * cfg.rankK;

    TPipe pipe;

    Hccl hccl;
    GM_ADDR contextGM = GetHcclContext<HCCL_GROUP_ID_0>();
    hccl.Init(contextGM, mc2InitTiling);
    hccl.SetCcTiling(mc2CcTiling);

    // 下发allgather任务
    // 首块
    auto handleId = hccl.AllGather<true>(aGM, gatherOutGM, aTileCnt, HcclDataType(cfg.dataType), aRankCnt, tileNum);
    // 尾块
    auto tailHandleId = hccl.AllGather<true>(aGM + tileNum * aTileOffset, gatherOutGM + tileNum * aTileOffset, aTailCnt,
        HcclDataType(cfg.dataType), aRankCnt, tailNum);

    if (TILING_KEY_IS(100)) { // full mesh + no nd2nz + biasNoNeedCast
        using A_TYPE = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, A_DTYPE, true>;
        using B_TYPE = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, B_DTYPE, true>;
        using C_TYPE = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, C_DTYPE>;
        using BIAS_TYPE = MatmulType<AscendC::TPosition::GM, CubeFormat::ND, typename BiasType<BIAS_DTYPE>::type>;

        // 本卡数据计算
        MatmulKernelLocal<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE>(aGM, bGM, biasGM, cGM, cfg, localTiling, hccl);

        // tile首块计算
        auto aAddr = gatherOutGM; // gatherOut 作为 mm A矩阵地址
        auto cAddr = cGM;
        if (tileNum > 0) {
            MatmulKernel<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE>(aAddr, bGM, biasGM, cAddr, cfg, tileTiling, hccl, handleId,
                tileNum);
        }

        // tail尾块计算
        aAddr = gatherOutGM + tileNum * aTileOffset;
        cAddr = cGM + tileNum * cTileOffset;
        if (tailNum > 0) {
            MatmulKernel<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE>(aAddr, bGM, biasGM, cAddr, cfg, tailTiling, hccl, tailHandleId,
                tailNum);
        }
    }
    hccl.Finalize();
}